package lib

func init() {
	Probs = append(Probs, Problem{
		Num:         295,
		Discription: "数据流中的中位数",
		Level:       3,
		Labels: map[string]int{
			"大小堆": 1,
			"奇偶性": 1,
		},
	})
}

//从上到下调整堆：插入时，取出堆顶给另一半的堆后，把堆底元素换到堆顶调整
//从下到上调整堆：插入堆
type MedianFinder struct {
	// 左边较小的那一半的大顶堆
	s1 []int
	//右边较大那一半的小顶堆
	s2 []int
}

func Constructor5() MedianFinder {
	return MedianFinder{}
}

//长度相等加左边，不等加右边
func (finder *MedianFinder) AddNum(num int) {
	if len(finder.s1) == len(finder.s2) {
		finder.insertSmallerHeap(num)
		smallest := finder.s2[0]
		finder.s2[0], finder.s2[len(finder.s2)-1] = finder.s2[len(finder.s2)-1], finder.s2[0]
		finder.s2 = finder.s2[:len(finder.s2)-1]
		finder.adjustSmallHeap(finder.s2, 0, len(finder.s2)-1)
		finder.insertBiggerHeap(smallest)
	} else {
		finder.insertBiggerHeap(num)
		biggest := finder.s1[0]
		finder.s1[0], finder.s1[len(finder.s1)-1] = finder.s1[len(finder.s1)-1], finder.s1[0]
		finder.s1 = finder.s1[:len(finder.s1)-1]
		finder.adjustBigHeap(finder.s1, 0, len(finder.s1)-1)
		finder.insertSmallerHeap(biggest)
	}
}

func (finder *MedianFinder) FindMedian() float64 {
	if len(finder.s1) == len(finder.s2) {
		return float64(finder.s1[0]+finder.s2[0]) / float64(2)
	} else {
		return float64(finder.s1[0])
	}
}

func (finder *MedianFinder) insertBiggerHeap(val int) {
	finder.s1 = append(finder.s1, val)
	i := len(finder.s1) - 1
	for (i-1)/2 >= 0 {
		parentIndex := (i - 1) / 2
		if finder.s1[parentIndex] < val {
			finder.s1[parentIndex], finder.s1[i] = finder.s1[i], finder.s1[parentIndex]
			i = parentIndex
		} else {
			break
		}
	}
}

func (finder *MedianFinder) insertSmallerHeap(val int) {
	finder.s2 = append(finder.s2, val)
	i := len(finder.s2) - 1
	for (i-1)/2 >= 0 {
		parentIndex := (i - 1) / 2
		if finder.s2[parentIndex] > val {
			finder.s2[parentIndex], finder.s2[i] = finder.s2[i], finder.s2[parentIndex]
			i = parentIndex
		} else {
			break
		}
	}
}

func (finder *MedianFinder) adjustBigHeap(nums []int, start int, end int) {
	//注意刚开始可能会出现这种情况，直接return
	if start >= end {
		return
	}

	topValue := nums[start]
	i := start
	for 2*i+1 <= end {
		leftChildIndex := 2*i + 1
		rightChildIndex := 2*i + 2
		biggerChildIndex := leftChildIndex
		if rightChildIndex <= end && nums[rightChildIndex] > nums[leftChildIndex] {
			biggerChildIndex = rightChildIndex
		}

		if topValue >= nums[biggerChildIndex] {
			break
		}

		nums[i], nums[biggerChildIndex] = nums[biggerChildIndex], nums[i]
		i = biggerChildIndex
	}

	//nums[i] = topValue
}

func (finder *MedianFinder) adjustSmallHeap(nums []int, start int, end int) {
	if start >= end {
		return
	}

	topValue := nums[start]
	i := start
	for 2*i+1 <= end {
		leftChildIndex := 2*i + 1
		rightChildIndex := 2*i + 2
		smallerChildIndex := leftChildIndex
		if rightChildIndex <= end && nums[rightChildIndex] < nums[leftChildIndex] {
			smallerChildIndex = rightChildIndex
		}

		if topValue <= nums[smallerChildIndex] {
			break
		}

		nums[i], nums[smallerChildIndex] = nums[smallerChildIndex], nums[i]
		i = smallerChildIndex
	}

	//nums[i] = topValue
}
